{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import wandb\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "from pprint import pprint\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib\n",
    "\n",
    "matplotlib.rcParams['mathtext.fontset'] = 'stix'\n",
    "matplotlib.rcParams['font.family'] = 'STIXGeneral'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sns.set_palette(sns.color_palette(\"Paired\"))\n",
    "matplotlib.rcParams.update({'font.size': 10})\n",
    "\n",
    "fig, axs = plt.subplots(1, 2, figsize=(12,3))\n",
    "locs = np.arange(4)\n",
    "\n",
    "heights = [0.2, 0.3, 0.95, 1.]\n",
    "\n",
    "bar1 = axs[0].bar(locs[:-1], heights[:-1], width=0.4, color=\"w\")\n",
    "bar1 += axs[0].bar(locs[-1:], heights[-1:], width=0.4, hatch=\"/\", color=\"w\")\n",
    "labels = [\"+ ODE Bias\", \"+ Second-Order\\nBias\", \"+ Symplectic\\nBias\", \"Hamiltonian\\nNN\"]\n",
    "for idx, (rect, label) in enumerate(zip(bar1, labels)):\n",
    "    height = rect.get_height()\n",
    "    if idx == (len(labels) - 1):\n",
    "        axs[0].text(rect.get_x() + rect.get_width() / 2.0, height + 0.09, label, ha='center', va='bottom',\n",
    "                    bbox=dict(facecolor='none', edgecolor='black', boxstyle='round,pad=0.5', linewidth=1))       \n",
    "    else:\n",
    "        axs[0].text(rect.get_x() + rect.get_width() / 2.0, height + 0.07, label, ha='center', va='bottom'\n",
    "\n",
    "axs[0].bar(locs[3:], heights[3], width=0.4) \n",
    "axs[0].bar(locs[2:], heights[2], width=0.4)\n",
    "axs[0].bar(locs[1:], heights[1], width=0.4) \n",
    "axs[0].bar(locs, heights[0], width=0.4)\n",
    "axs[0].bar(locs, heights, width=0.4, edgecolor='black', fill=False) \n",
    "    \n",
    "axs[0].set_title(\"Expectation\", fontsize=20, pad=40)\n",
    "axs[0].set_yticks(np.linspace(0, 1, 5))\n",
    "axs[0].set_yticklabels([\"\", \"25%\", \"50%\", \"75%\", \"100%\"], fontsize=14)\n",
    "axs[0].set_ylabel(\"Performance\", fontsize=16)#, rotation=45)\n",
    "axs[0].tick_params(axis=u'both', which=u'both',length=0)\n",
    "\n",
    "axs[0].set_xticks([])\n",
    "axs[0].spines['top'].set_visible(False)\n",
    "axs[0].spines['right'].set_visible(False)\n",
    "\n",
    "heights = [0.2, 0.9, 0.95, 1.0]\n",
    "\n",
    "bar2 = axs[1].bar(locs[:-1], heights[:-1], width=0.4, color=\"w\", edgecolor='black') \n",
    "bar2 += axs[1].bar(locs[-1:], heights[-1:], width=0.4, hatch=\"/\", color=\"w\", edgecolor='black')\n",
    "# labels = [\"+ ODE Bias\", \"+ Second-Order\\nBias\", \"+ Symplectic\\nBias\", \"+ Simple Function\\nBias\", \"HNN\"]\n",
    "labels = [\"+ ODE Bias\", \"+ Second-Order\\nBias\", \"+ Symplectic\\nBias\", \"Hamiltonian\\nNN\"]\n",
    "for idx, (rect, label) in enumerate(zip(bar2, labels)):\n",
    "    height = rect.get_height()\n",
    "    if idx == (len(labels) - 1):\n",
    "        axs[1].text(rect.get_x() + rect.get_width() / 2.0, height + 0.09, label, ha='center', va='bottom',\n",
    "                    bbox=dict(facecolor='none', edgecolor='black', boxstyle='round,pad=0.5', linewidth=1))       \n",
    "    else:\n",
    "        axs[1].text(rect.get_x() + rect.get_width() / 2.0, height + 0.07, label, ha='center', va='bottom')\n",
    "\n",
    "axs[1].bar(locs[3:], heights[3], width=0.4)\n",
    "axs[1].bar(locs[2:], heights[2], width=0.4)\n",
    "axs[1].bar(locs[1:], heights[1], width=0.4)\n",
    "axs[1].bar(locs, heights[0], width=0.4)\n",
    "axs[1].bar(locs, heights, width=0.4, edgecolor='black', fill=False) \n",
    "    \n",
    "axs[1].set_title(\"Reality\", fontsize=20, pad=40)\n",
    "axs[1].set_yticks(np.linspace(0, 1, 5))\n",
    "axs[1].set_yticklabels([\"\", \"25%\", \"50%\", \"75%\", \"100%\"], fontsize=14)\n",
    "axs[1].set_ylabel(\"Performance\", fontsize=16)\n",
    "axs[1].tick_params(axis=u'both', which=u'both',length=0)\n",
    "\n",
    "axs[1].set_xticks([])\n",
    "axs[1].spines['top'].set_visible(False)\n",
    "axs[1].spines['right'].set_visible(False)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "fig.savefig(\"figure1.pdf\", bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "download_root = \".\"\n",
    "\n",
    "def get_sweep_regression_df_all(sweep_id, allow_crash=False):\n",
    "    api = wandb.Api()\n",
    "    sweep = api.sweep(\"ngruver/physics-uncertainty-exps/{}\".format(sweep_id))\n",
    "    \n",
    "    results = []\n",
    "    for run in sweep.runs:        \n",
    "        config = pd.Series(run.config)\n",
    "        \n",
    "        if not allow_crash and not \"finished\" in str(run):\n",
    "            continue\n",
    "        \n",
    "        if \"finished\" in str(run):\n",
    "            summary = pd.Series(run.summary)\n",
    "        else:\n",
    "            history = run.history()\n",
    "            summary = pd.Series({k: history[k].to_numpy()[-1] for k,v in history.items()})\n",
    "        results.append(pd.concat([config,summary]))\n",
    "    return pd.concat(results,axis=1).T\n",
    "\n",
    "\n",
    "sweep_id = \"v96kirjy\"\n",
    "df = get_sweep_regression_df_all(sweep_id, allow_crash=True)\n",
    "sweep_id = \"pexiwka8\"\n",
    "df_2 = get_sweep_regression_df_all(sweep_id, allow_crash=True)\n",
    "full_df = pd.concat([df, df_2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = full_df[full_df[\"hidden_size\"].isin([256])]\n",
    "\n",
    "metric = \"test_gerr\"\n",
    "_df = df[df[\"num_bodies\"].isin([2,3,4])]\n",
    "gb = _df.groupby([\"model_type\", \"system_type\", \"num_bodies\"])[metric].apply(list).reset_index(name=\"vals\")\n",
    "\n",
    "model_types = [[\"NN\", \"MechanicsNN\"],\n",
    "               [\"NonseparableHNN\", \"HNN\"]]\n",
    "titles = [\"NODE Models\", \"HNN Models\"]\n",
    "groups = [\"Without SO bias\", \"With SO bias\"]\n",
    "\n",
    "colors = [\"#6A0078\", \"#dd63ff\"]\n",
    "sns.set_palette(sns.color_palette(\"Paired\", desat=0.8)[2:4])\n",
    "matplotlib.rcParams.update({'font.size': 14})\n",
    "\n",
    "gap_width = 0.4\n",
    "fig, axs = plt.subplots(1, 2, figsize=(12,3), sharey=True)\n",
    "\n",
    "for i in range(len(model_types)):\n",
    "    pts = gb[gb[\"model_type\"].isin(model_types[i])]\n",
    "    \n",
    "    num_gaps = (len(pts) // len(model_types[i])) - 1\n",
    "    num_locs = len(pts) + gap_width * num_gaps\n",
    "    interval = len(model_types[i]) + gap_width\n",
    "    \n",
    "    for j in range(len(model_types[0])):\n",
    "        locs = np.arange(j, num_locs, interval)[:num_gaps+1]\n",
    "        _pts = pts[pts[\"model_type\"] == model_types[i][j]]\n",
    "        model_vals = _pts[\"vals\"]\n",
    "        system_types = [x.replace('Pendulum','') for x in _pts[\"system_type\"]]\n",
    "        num_bodies = _pts[\"num_bodies\"]\n",
    "\n",
    "        means = [np.mean(x) for x in model_vals]\n",
    "        errs = [np.std(x) / len(x) for x in model_vals]\n",
    "        \n",
    "        axs[i].bar(locs, means, width=0.95, yerr=errs, label=groups[j])    \n",
    "        axs[i].set_yscale('log')\n",
    "        axs[i].set_yticks([], minor=True)\n",
    "        if i == 0:\n",
    "            axs[i].set_ylabel(\"Rollout Error\")\n",
    "\n",
    "        axs[i].set_xticks(np.arange(0.5, num_locs, interval))\n",
    "        axs[i].set_xticklabels([f\"{system} {n}\" for system, n in zip(system_types, num_bodies)])\n",
    "\n",
    "        axs[i].tick_params(axis=u'both', which=u'both',length=0)\n",
    "        \n",
    "        axs[i].set_title(titles[i], fontsize=20, pad=10)\n",
    "    \n",
    "plt.tight_layout()\n",
    "    \n",
    "handles, labels = axs[0].get_legend_handles_labels()\n",
    "fig.legend(handles, labels, loc='lower center', ncol=3)\n",
    "fig.subplots_adjust(bottom=0.27)\n",
    "\n",
    "plt.show()\n",
    "\n",
    "fig.savefig(\"figure3.pdf\", bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = full_df[full_df[\"hidden_size\"].isin([256])]\n",
    "\n",
    "metric = \"test_MSE\"\n",
    "_df = df[df[\"num_bodies\"].isin([1,2,3,4])]\n",
    "gb = _df.groupby([\"model_type\", \"system_type\", \"num_bodies\"])[metric].apply(list).reset_index(name=\"vals\")\n",
    "\n",
    "model_types = [[\"MechanicsNN\", \"HNN\"],\n",
    "               [\"MechanicsNN\", \"HNN\"]]\n",
    "systems = [\"ChainPendulum\", \"SpringPendulum\"]\n",
    "titles = [\"Complex Coordinates\", \"Simple Coordinates\"]\n",
    "groups = [\"NODE + SO\", \"HNN\"]\n",
    "\n",
    "colors = [\"#00058A\", \"#6A0078\"]\n",
    "sns.set_palette(sns.color_palette(colors))\n",
    "matplotlib.rcParams.update({'font.size': 14})\n",
    "\n",
    "gap_width = 1.0\n",
    "fig, axs = plt.subplots(1, 2, figsize=(12,3))\n",
    "\n",
    "for i in range(len(model_types)):\n",
    "    pts = gb[(gb[\"model_type\"].isin(model_types[i])) & (gb[\"system_type\"] == systems[i])]\n",
    "    \n",
    "    num_gaps = (len(pts) // len(model_types[i])) - 1\n",
    "    num_locs = len(pts) + gap_width * num_gaps\n",
    "    interval = len(model_types[i]) + gap_width\n",
    "    \n",
    "    for j in range(len(model_types[0])):\n",
    "        locs = np.arange(j, num_locs, interval)[:num_gaps+1]\n",
    "        _pts = pts[pts[\"model_type\"] == model_types[i][j]]\n",
    "        model_vals = _pts[\"vals\"]\n",
    "        system_types = [x.replace('Pendulum','') for x in _pts[\"system_type\"]]\n",
    "        num_bodies = _pts[\"num_bodies\"]\n",
    "\n",
    "        means = [np.mean(x) for x in model_vals]\n",
    "        errs = [np.std(x) / len(x) for x in model_vals]\n",
    "        \n",
    "        axs[i].bar(locs, means, width=0.95, yerr=errs, label=groups[j])    \n",
    "        axs[i].set_yscale('log')\n",
    "        axs[i].set_yticks([], minor=True)\n",
    "        axs[i].set_ylabel(\"Rollout Error\")\n",
    "\n",
    "        axs[i].set_xticks(np.arange(0.5, num_locs, interval))\n",
    "        axs[i].set_xticklabels([f\"{system} {n}\" for system, n in zip(system_types, num_bodies)])\n",
    "\n",
    "        axs[i].tick_params(axis=u'both', which=u'both',length=0)\n",
    "        \n",
    "        axs[i].set_title(titles[i], fontsize=20, pad=10)\n",
    "    \n",
    "plt.tight_layout()\n",
    "    \n",
    "handles, labels = axs[0].get_legend_handles_labels()\n",
    "fig.legend(handles, labels, loc='lower center', ncol=3)\n",
    "fig.subplots_adjust(bottom=0.27)\n",
    "\n",
    "plt.show()\n",
    "\n",
    "fig.savefig(\"figure4.pdf\", bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sweep_id = \"vp3sn5gg\"\n",
    "friction_df = get_sweep_regression_df_all(sweep_id, allow_crash=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = friction_df\n",
    "\n",
    "metric = \"test_gerr\"\n",
    "_df = df[df[\"num_bodies\"].isin([1,2,3,4])]\n",
    "gb = _df.groupby([\"model_type\", \"system_type\", \"num_bodies\"])[metric].apply(list).reset_index(name=\"vals\")\n",
    "\n",
    "model_types = [\"NN\", \"MechanicsNN\", \"HNN\", \"MixtureHNN\"]\n",
    "reader_friendly_dict = {\n",
    "    \"NN\": \"NODE\", \n",
    "    \"MechanicsNN\": \"NODE + SO\",\n",
    "    \"HNN\": \"HNN\",\n",
    "    \"MixtureHNN\": \"SymODEN\"\n",
    "}\n",
    "\n",
    "colors = [\"#00abdf\", \"#00058A\", \"#6A0078\", (96/255,74/255,123/255), \"#8E6100\"]\n",
    "sns.set_palette(sns.color_palette(colors))\n",
    "\n",
    "gap_width = 1.\n",
    "fig, ax = plt.subplots(1, 1, figsize=(10,3))\n",
    "\n",
    "pts = gb[gb[\"model_type\"].isin(model_types)]\n",
    "\n",
    "num_gaps = (len(pts) // len(model_types)) - 1\n",
    "num_locs = len(pts) + gap_width * num_gaps\n",
    "interval = len(model_types) + gap_width\n",
    "\n",
    "for i in range(len(model_types)):\n",
    "    locs = np.arange(i, num_locs, interval)[:num_gaps+1]\n",
    "    _pts = pts[pts[\"model_type\"] == model_types[i]]\n",
    "\n",
    "    model_vals = _pts[\"vals\"]\n",
    "    system_types = [x.replace('Pendulum','').replace('Friction','') for x in _pts[\"system_type\"]]\n",
    "    num_bodies = _pts[\"num_bodies\"]\n",
    "\n",
    "    means = [np.mean([v for v in x if not np.isnan(v)]) for x in model_vals]\n",
    "    errs = [np.std([v for v in x if not np.isnan(v)]) / len(x) for x in model_vals]\n",
    "\n",
    "    ax.bar(locs, means, width=0.95, yerr=errs, label=reader_friendly_dict[model_types[i]])    \n",
    "\n",
    "ax.set_yscale('log')\n",
    "ax.set_yticks([], minor=True)\n",
    "ax.set_ylabel(\"Rollout Error\")\n",
    "\n",
    "ax.set_xticks(np.arange(1.5, num_locs, interval))\n",
    "ax.set_xticklabels([f\"{system} {n}\" for system, n in zip(system_types, num_bodies)])\n",
    "\n",
    "ax.tick_params(axis=u'both', which=u'both',length=0)\n",
    "\n",
    "ax.set_xlim(-1, num_locs)\n",
    "# ax.set_title(titles[i])\n",
    "    \n",
    "plt.tight_layout()\n",
    "    \n",
    "handles, labels = ax.get_legend_handles_labels()\n",
    "fig.legend(handles, labels, loc='lower center', ncol=4)\n",
    "fig.subplots_adjust(bottom=0.27)\n",
    "\n",
    "plt.show()\n",
    "\n",
    "fig.savefig(\"figure5.pdf\", bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_sweep_regression_df_all(sweep_id, allow_crash=False):\n",
    "    api = wandb.Api()\n",
    "    sweep = api.sweep(\"samuelstanton/physics-uncertainty-exps/{}\".format(sweep_id))\n",
    "    \n",
    "    results = []\n",
    "    for run in sweep.runs:        \n",
    "        config = pd.Series(run.config)\n",
    "        \n",
    "        if not allow_crash and not \"finished\" in str(run):\n",
    "            continue\n",
    "        \n",
    "        if \"finished\" in str(run):\n",
    "            summary = pd.Series(run.summary)\n",
    "        else:\n",
    "            history = run.history()\n",
    "            summary = pd.Series({k: history[k].to_numpy()[-1] for k,v in history.items()})\n",
    "#         print(summary)\n",
    "        results.append(pd.concat([config,summary]))\n",
    "    return pd.concat(results,axis=1).T\n",
    "\n",
    "\n",
    "sweep_id = \"wkn77qza\"\n",
    "df = get_sweep_regression_df_all(sweep_id, allow_crash=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sweep_id = \"ahybaa2m\"\n",
    "df_2 = get_sweep_regression_df_all(sweep_id, allow_crash=True)\n",
    "df = pd.concat([df, df_2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sweep_id = \"glupfh56\"\n",
    "df_3 = get_sweep_regression_df_all(sweep_id, allow_crash=True)\n",
    "df = pd.concat([df, df_3])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = df[(df[\"hidden_size\"] == 128) & (df[\"subsample_ratio\"] == 1.)]\n",
    "df[\"task\"] = df[\"task\"].apply(lambda s: s.replace(\"Full-v0\", \"\"))\n",
    "\n",
    "reader_friendly_dict = {\n",
    "    \"NODE\": \"NODE\",\n",
    "    \"CoupledNODE\": \"NODE + SO\",\n",
    "    \"MixtureHNN\": \"SymODEN\"\n",
    "}\n",
    "\n",
    "def bar_plot(ax, title, df, metric, group_label, methods, ind_label, keys=None):\n",
    "    if keys is None:\n",
    "        keys = df[ind_label].drop_duplicates().to_numpy()\n",
    "        keys = keys[np.argsort(keys)]\n",
    "    \n",
    "    for idx, method in enumerate(methods):\n",
    "        _df = df[df[group_label] == method]\n",
    "        locs = idx + (len(methods) + 1) * np.arange(len(keys))\n",
    "        means, stds = [],[]\n",
    "        for idx, k in enumerate(keys):\n",
    "            try:\n",
    "                __df = _df[_df[ind_label] == k]\n",
    "                __df = __df[~__df[metric].isna()]\n",
    "                if title == \"Test\":\n",
    "                    data = [x for x in __df[metric].to_numpy()]\n",
    "                else:\n",
    "                    data = [x for x in __df[metric].to_numpy()]\n",
    "                means.append(np.mean(data))\n",
    "                stds.append(np.std(data))\n",
    "            except Exception as e:\n",
    "                print(e)\n",
    "                locs = np.delete(locs, idx)\n",
    "                continue\n",
    "        label = method if method != \"FlexHNN\" else \"RPP\"\n",
    "        ax.bar(locs, means, yerr=stds, label=reader_friendly_dict[method], width=0.9)\n",
    "    \n",
    "    ax.set_yscale('log')\n",
    "    ax.set_xticks(np.arange(1, len(keys) * len(methods) + len(keys) - 1, len(methods) + 1))\n",
    "    ax.set_xticklabels(keys)#, rotation = 20)\n",
    "    ax.set_ylabel(\"Rollout Err\")\n",
    "    ax.tick_params(axis=u'both', which=u'both',length=0)\n",
    "    ax.set_title(title, fontsize=22, pad=10) \n",
    "\n",
    "colors = [\"#00abdf\", \"#00058A\", (96/255,74/255,123/255), \"#8E6100\"]\n",
    "sns.set_palette(sns.color_palette(colors))\n",
    "matplotlib.rcParams.update({'font.size': 18})\n",
    "\n",
    "metric = \"test_mse\"\n",
    "methods = [\"NODE\", \"CoupledNODE\", \"MixtureHNN\"]\n",
    "\n",
    "fig, ax = plt.subplots(1, 2, figsize=(12,4))\n",
    "bar_plot(ax[0], \"Test\", df, \"test_mse\", \"model_type\", methods, \"task\")\n",
    "bar_plot(ax[1], \"Train\", df, \"train_mse\", \"model_type\", methods, \"task\")\n",
    "\n",
    "plt.tight_layout()\n",
    "\n",
    "handles, labels = ax[0].get_legend_handles_labels()\n",
    "fig.legend(handles, labels, ncol=3, loc='lower center')\n",
    "fig.subplots_adjust(bottom=0.25)#, left=-0.4)\n",
    "\n",
    "plt.show()\n",
    "\n",
    "fig.savefig(\"figure7.pdf\", bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_sweep_regression_df_all(sweep_id, allow_crash=False):\n",
    "    api = wandb.Api()\n",
    "    sweep = api.sweep(\"ngruver/physics-uncertainty-exps/{}\".format(sweep_id))\n",
    "    \n",
    "    results = []\n",
    "    for run in sweep.runs:        \n",
    "        config = pd.Series(run.config)\n",
    "        \n",
    "        if not allow_crash and not \"finished\" in str(run):\n",
    "            continue\n",
    "        \n",
    "        if \"finished\" in str(run):\n",
    "            summary = pd.Series(run.summary)\n",
    "        else:\n",
    "            history = run.history()\n",
    "            summary = pd.Series({k: history[k].to_numpy()[-1] for k,v in history.items()})\n",
    "        results.append(pd.concat([config,summary]))\n",
    "    return pd.concat(results,axis=1).T\n",
    "\n",
    "sweep_id = \"9f0ldbr6\"\n",
    "df1 = get_sweep_regression_df_all(sweep_id, allow_crash=True)\n",
    "sweep_id = \"p400ew2f\"\n",
    "df2 = get_sweep_regression_df_all(sweep_id, allow_crash=True)\n",
    "df = pd.concat([df1, df2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "metric = \"test_gerr\"\n",
    "_df = df[df[\"num_bodies\"].isin([1,2,3,4])]\n",
    "gb = _df.groupby([\"model_type\", \"system_type\", \"loss\", \"hidden_size\"])[metric].apply(list).reset_index(name=\"vals\")\n",
    "gb[\"mean\"] = pd.DataFrame(gb[\"vals\"].values.tolist()).mean(1)\n",
    "gb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gb = gb[gb[\"hidden_size\"] == 256]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_types = [\"NN\", \"MechanicsNN\", \"HNN\"]\n",
    "reader_friendly_dict = {\n",
    "    \"NN\": \"NODE\", \n",
    "    \"MechanicsNN\": \"NODE + SO\",\n",
    "    \"HNN\": \"HNN\",\n",
    "}\n",
    "\n",
    "colors = [\"#00abdf\", \"#00058A\", \"#6A0078\", (96/255,74/255,123/255), \"#8E6100\"]\n",
    "sns.set_palette(sns.color_palette(colors))\n",
    "matplotlib.rcParams.update({'font.size': 18})\n",
    "\n",
    "gap_width = 1.\n",
    "fig, ax = plt.subplots(1, 1, figsize=(10,3))\n",
    "\n",
    "pts = gb[gb[\"model_type\"].isin(model_types)]\n",
    "\n",
    "num_gaps = (len(pts) // len(model_types)) - 1\n",
    "num_locs = len(pts) + gap_width * num_gaps\n",
    "interval = len(model_types) + gap_width\n",
    "\n",
    "for i in range(len(model_types)):\n",
    "    locs = np.arange(i, num_locs, interval)[:num_gaps+1]\n",
    "    _pts = pts[pts[\"model_type\"] == model_types[i]]\n",
    "\n",
    "    model_vals = _pts[\"vals\"]\n",
    "    system_types = [x.replace('Pendulum','').replace('Friction','') for x in _pts[\"system_type\"]]\n",
    "    num_bodies = _pts[\"loss\"]\n",
    "    hidden_size = _pts[\"hidden_size\"]\n",
    "\n",
    "    means = [np.mean([v for v in x if not np.isnan(v)]) for x in model_vals]\n",
    "    errs = [np.std([v for v in x if not np.isnan(v)]) / len(x) for x in model_vals]\n",
    "\n",
    "    ax.bar(locs, means, width=0.95, yerr=errs, label=reader_friendly_dict[model_types[i]])    \n",
    "\n",
    "ax.set_yscale('log')\n",
    "ax.set_yticks([], minor=True)\n",
    "ax.set_ylabel(\"Rollout Error\")\n",
    "\n",
    "ax.set_xticks(np.arange(1.5, num_locs, interval))\n",
    "ax.set_xticklabels([f\"{system} ({n} loss)\" for system, n in zip(system_types, num_bodies)])\n",
    "ax.tick_params(axis=u'both', which=u'both',length=0)\n",
    "\n",
    "ax.set_xlim(-1, num_locs)\n",
    " \n",
    "plt.tight_layout()\n",
    "    \n",
    "handles, labels = ax.get_legend_handles_labels()\n",
    "fig.legend(handles, labels, loc='lower center', ncol=4)\n",
    "fig.subplots_adjust(bottom=0.35)\n",
    "\n",
    "plt.show()\n",
    "\n",
    "fig.savefig(\"l1-l2-loss-comparison.pdf\", bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tqdm \n",
    "\n",
    "def get_sweep_regression_df_all(sweep_id, allow_crash=False):\n",
    "    api = wandb.Api(timeout=60)\n",
    "    sweep = api.sweep(\"ngruver/physics-uncertainty-exps/{}\".format(sweep_id))\n",
    "    \n",
    "    results = []\n",
    "    for run in tqdm.tqdm(sweep.runs):        \n",
    "        config = pd.Series(run.config)\n",
    "        \n",
    "        if not allow_crash and not \"finished\" in str(run):\n",
    "            continue\n",
    "             \n",
    "        history = run.history()\n",
    "        summary = pd.Series({k: min(history[k].to_numpy()) for k,v in history.items()})\n",
    "        results.append(pd.concat([config,summary]))\n",
    "    return pd.concat(results,axis=1).T\n",
    "\n",
    "sweep_id = \"s9fjv7t1\" #\"/tpijuz7e\"\n",
    "df = get_sweep_regression_df_all(sweep_id, allow_crash=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = _df[_df[\"n_systems\"] == 5000]\n",
    "metric = \"test_gerr\"\n",
    "gb = df.groupby([\"model_type\", \"system_type\", \"n_systems\", \"data_seed\"])[metric].apply(list).reset_index(name=\"vals\")\n",
    "gb[\"min_val\"] = pd.DataFrame(gb[\"vals\"].values.tolist()).min(1)\n",
    "gb = gb.drop(columns=['vals'])\n",
    "gb = gb.groupby([\"model_type\", \"system_type\", \"n_systems\"])[\"min_val\"].apply(list).reset_index(name=\"vals\")\n",
    "gb[\"mean\"] = pd.DataFrame(gb[\"vals\"].values.tolist()).mean(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_types = [\"NN\", \"MechanicsNN\", \"HNN\"]\n",
    "reader_friendly_dict = {\n",
    "    \"NN\": \"NODE\", \n",
    "    \"MechanicsNN\": \"NODE + SO\",\n",
    "    \"HNN\": \"HNN\",\n",
    "}\n",
    "\n",
    "colors = [\"#00abdf\", \"#00058A\", \"#6A0078\", (96/255,74/255,123/255), \"#8E6100\"]\n",
    "sns.set_palette(sns.color_palette(colors))\n",
    "matplotlib.rcParams.update({'font.size': 18})\n",
    "\n",
    "gap_width = 1.\n",
    "fig, ax = plt.subplots(1, 1, figsize=(10,3))\n",
    "\n",
    "pts = gb[gb[\"model_type\"].isin(model_types)]\n",
    "\n",
    "num_gaps = (len(pts) // len(model_types)) - 1\n",
    "num_locs = len(pts) + gap_width * num_gaps\n",
    "interval = len(model_types) + gap_width\n",
    "\n",
    "for i in range(len(model_types)):\n",
    "    locs = np.arange(i, num_locs, interval)[:num_gaps+1]\n",
    "    _pts = pts[pts[\"model_type\"] == model_types[i]]\n",
    "\n",
    "    print(_pts.drop(columns=['vals']))\n",
    "    model_vals = _pts[\"vals\"]\n",
    "    system_types = [x.replace('Pendulum','').replace('Friction','') for x in _pts[\"system_type\"]]\n",
    "    n_systems = _pts[\"n_systems\"]\n",
    "    \n",
    "    means = [np.mean([v for v in x if not np.isnan(v)]) for x in model_vals]\n",
    "    errs = [np.std([v for v in x if not np.isnan(v)]) / len(x) for x in model_vals]\n",
    "\n",
    "    ax.bar(locs, means, width=0.95, yerr=errs, label=reader_friendly_dict[model_types[i]])    \n",
    "\n",
    "ax.set_yscale('log')\n",
    "ax.set_yticks([], minor=True)\n",
    "ax.set_ylabel(\"Rollout Error\")\n",
    "\n",
    "ax.set_xticks(np.arange(1.5, num_locs, interval))\n",
    "ax.set_xticklabels([f\"{system}\" for n, system in zip(n_systems, system_types)])\n",
    "\n",
    "ax.tick_params(axis=u'both', which=u'both',length=0)\n",
    "\n",
    "ax.set_xlim(-1, num_locs)\n",
    "   \n",
    "plt.tight_layout()\n",
    "    \n",
    "handles, labels = ax.get_legend_handles_labels()\n",
    "fig.legend(handles, labels, loc='lower center', ncol=4)\n",
    "fig.subplots_adjust(bottom=0.35)\n",
    "\n",
    "plt.show()\n",
    "\n",
    "fig.savefig(\"additional_systems.pdf\", bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import ticker\n",
    "\n",
    "colors = [\"#00abdf\", \"#00058A\", \"#6A0078\", (96/255,74/255,123/255), \"#8E6100\"]\n",
    "sns.set_palette(sns.color_palette(colors))\n",
    "matplotlib.rcParams.update({'font.size': 18})\n",
    "\n",
    "fig, ax = plt.subplots(1, 1, figsize=(3,3))\n",
    "\n",
    "ax.bar(0, 89.8190, width=0.95, label=\"NODE\")    \n",
    "ax.bar(1, 85.6026, width=0.95, label=\"NODE + SO\") \n",
    "ax.bar(2, 85.1126, width=0.95, label=\"HNN\") \n",
    "\n",
    "ax.set_yscale('log')\n",
    "ax.set_yticks([86, 88, 90], minor=True)\n",
    "ax.yaxis.set_minor_formatter(ticker.ScalarFormatter())\n",
    "ax.set_ylabel(\"Rollout Error\")\n",
    "\n",
    "ax.tick_params(axis=u'both', which=u'both',length=0)\n",
    "\n",
    "ax.set_xlim(-1, 3)\n",
    "ax.set_xticks([1])\n",
    "ax.set_xticklabels([\"Pendulum w/ Contact\"])\n",
    "    \n",
    "plt.tight_layout()\n",
    "    \n",
    "handles, labels = ax.get_legend_handles_labels()\n",
    "fig.legend(handles, labels, loc='lower center', ncol=3)\n",
    "fig.subplots_adjust(bottom=0.35)\n",
    "\n",
    "plt.show()\n",
    "\n",
    "fig.savefig(\"contact_system.pdf\", bbox_inches='tight')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
